import os
import joblib
import numpy as np
from skimage import io, color
from skimage.transform import resize
from skimage.feature import hog
from sklearn import svm
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import LabelEncoder

# 图像预处理和特征提取
def extract_features(image_path, img_size=(1024, 691)):
    image = io.imread(image_path)
    if len(image.shape) == 3:  # 如果是彩色图像
        image = color.rgb2gray(image)
    image_resized = resize(image, img_size)
    features, _ = hog(image_resized, block_norm='L2-Hys', pixels_per_cell=(16, 16),
                      cells_per_block=(2, 2), visualize=True)
    return features

# 加载数据集
def load_dataset(dataset_path):
    features = []
    labels = []
    for label in os.listdir(dataset_path):
        class_path = os.path.join(dataset_path, label)
        if os.path.isdir(class_path):
            for img_name in os.listdir(class_path):
                img_path = os.path.join(class_path, img_name)
                if img_path.endswith('tif'):
                    features.append(extract_features(img_path))
                    labels.append(label)
    return np.array(features), np.array(labels)

# 数据集路径
train_dataset_path = r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\lab7_SVM\data\train'
test_dataset_path = r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\lab7_SVM\data\test'

# 加载数据
X_train, y_train = load_dataset(train_dataset_path)
X_test, y_test = load_dataset(test_dataset_path)

# 标签编码
label_encoder = LabelEncoder()
y_train_encoded = label_encoder.fit_transform(y_train)
y_test_encoded = label_encoder.transform(y_test)

# 保存标签编码器
joblib.dump(label_encoder, 'label_encoder.pkl')

# 定义SVM分类器
clf = svm.SVC(kernel='linear', probability=True)

# 初始化用于跟踪最佳模型的变量
best_accuracy = 0.0
num_epochs = 100

# 训练和测试模型
for epoch in range(num_epochs):
    # 训练SVM分类器
    clf.fit(X_train, y_train_encoded)

    # 测试分类器
    y_pred = clf.predict(X_test)
    accuracy = accuracy_score(y_test_encoded, y_pred)
    print(f'Epoch {epoch+1}/{num_epochs}, Accuracy: {accuracy * 100:.2f}%')

    # 保存最好的模型
    if accuracy >= best_accuracy:
        best_accuracy = accuracy
        joblib.dump(clf, 'svm_hog_best_model.pkl')
        print(f'Saved Best Model with Accuracy: {accuracy * 100:.2f}%')
